
#ifndef __neutronics_system_h__
#define __neutronics_system_h__

#define ENERGY_FLUXES "flux_*"

// C++ includes
#include <map>

// libmesh includes
#include "libmesh.h"
#include "equation_systems.h"
#include "explicit_system.h"
#include "enum_eigen_solver_type.h"
#include "petsc_vector.h"
#include "petsc_matrix.h"
#include "perf_log.h"

#include "petscksp.h"

#define PUSH_LOG(a, b) NeutronicsSystem::logger.push(a, b)
#define POP_LOG(a, b) NeutronicsSystem::logger.pop(a, b)
//#define PUSH_LOG(a) NeutronicsSystem::logger.push(a, "EigenSolver")
//#define POP_LOG(a) NeutronicsSystem::logger.pop(a, "EigenSolver")

// Forward declaration
class Problem ;
namespace libMesh
{
	class MeshBase ;
}

class NeutronicsSystem : public EquationSystems
{
public:

	/**
	* Constructor.  Optionally initializes required
	* data structures.
	*/
	NeutronicsSystem (  libMesh::MeshBase& o_mesh, const Problem&, 
						const unsigned int e_groups,
						const std::string order = "FIRST",
						const std::string family = "LAGRANGE") ;

	/**
	* Destructor.
	*/
	virtual ~NeutronicsSystem ();

	/**
	* Clear all the data structures associated with
	* the current Physics. 
	*/
	virtual void clear ();

	/**
	* Initializes all the data structures necessary for computation using
	* the current Physics.
	*/
	virtual void init ();
	
	/**
	* Re-Initializes the data structures when mesh changes.
	*/
	virtual void reinit ();
	
	Real solve_eigen_problem(Real tol=1e-4, Real maxit=10) ;
	
	// Problem specific methods

	// The order for variableparam coupled fields = 
	//	Fuel Temperature, Moderator Density, Boron ppm, Moderator temp
	// It is not necessary to have all the fields when calling a cross-section function ; 
	// If no fields are specified, the default values DEFAULT_FUEL_TEMP, DEFAULT_MOD_DEN are used

	Real Keff() const ;
	
	Real& Keff() ;

	bool is_eigen_problem() const ;
 
	bool is_eigen_problem_matrix_free() const ;
 	
	bool& recompute_group_blocks() ;
	
	bool recompute_group_blocks() const ;

	bool& only_diagonal_loss_matrix() ;
	
	bool only_diagonal_loss_matrix() const ;

	bool& only_block_diagonal_loss_matrix() ;
	
	bool only_block_diagonal_loss_matrix() const ;
	
	///////////////////////////////////////////////////////////////////////////////////////////
	////////////////////////// OPERATOR SPECIFIC ROUTINES /////////////////////////////////////
	///////////////////////////////////////////////////////////////////////////////////////////
	PetscErrorCode assemble_loss_matrix() ;
	
	PetscErrorCode assemble_gain_matrix() ;
	
	PetscErrorCode gain_matrix_vector(const PetscVector<Real>&, PetscVector<Real>&) ;
	
	PetscErrorCode loss_matrix_vector(const PetscVector<Real>&, PetscVector<Real>&) ;

	// prec_context control whether an exact or inexact solve is performed
	PetscErrorCode block_gauss_seidel_solve(const PetscVector<Real>&, PetscVector<Real>&, bool prec_context=false) ;
	////////////////////////////////////////////////////////////////////////////////////////////

	PetscMatrix<Real>* get_loss_matrix() const ;
	
	PetscMatrix<Real>* get_gain_matrix() const ;
	
	EigenProblemType eigen_problem_type() const ;
	
	double& eigen_shift() ;
	
	double eigen_shift() const ;
	
	KSP& loss_matrix_solver() ;
	
	void initialize_eigen_context(bool is_mfree=false, EigenProblemType eigen_type=GNHEP) ;
	
	const PetscVector<Real>& get_solution() const ;
	PetscVector<Real>& get_solution() ;
	
	const PetscVector<Real>& get_residual() const ;
	PetscVector<Real>& get_residual() ;
	
	PetscVector<Real>& group_vector(unsigned int ig) const ;
	
	PetscMatrix<Real>& group_matrix(unsigned int ig, unsigned int jg) const ;
	
	KSP& group_solver(unsigned int ig) const ;

	virtual PetscErrorCode form_initial_guess (PetscVector<Real>& xs) ;

	const Problem& get_problem() const ;
	
	// changes the default values : useful when MMS is linearized about a certain state-space
	unsigned int energy_groups() const ;
	
	std::vector<unsigned int> flx_var ;
	
	static PerfLog logger ;

private:

	// Physics specific properties and routines
	const Problem* _problem ;

	// number of energy groups
	const unsigned int _egroups ;

	bool _initialized ;
	 
	// eigenvalue related members
	bool _is_eigen_problem;
	bool _is_gain_matrix_created, _is_loss_matrix_created ;
	EigenProblemType _eigen_problem_type ;
	
	bool _is_eigen_problem_mfree ;

	bool _only_block_diagonal_loss_matrix ;
	
	bool _only_diagonal_loss_matrix ;
	
	Real _keff ;
	Real _eigen_shift ;
	PetscMatrix<Real> *_loss_mat, *_gain_mat ;
	KSP _loss_matrix_solver ;

	ExplicitSystem& _system ;
	
	// preconditioning helpers - group block matrices, solvers
	KSP* _group_block_solvers ;
	PetscMatrix<Real>** _group_block_matrices ;
	PetscVector<Real>** _group_block_vectors ;
	
	bool _recompute_group_blocks ;
	
};

// Inline functions for NeutronicsSystem

inline
Real NeutronicsSystem::Keff() const
{
	return _keff ;
}

inline
Real& NeutronicsSystem::Keff() 
{
	return _keff ;
}

inline
unsigned int NeutronicsSystem::energy_groups() const
{
	return _egroups ;
}

inline
bool NeutronicsSystem::is_eigen_problem() const
{
	return _is_eigen_problem ;
}

inline
EigenProblemType NeutronicsSystem::eigen_problem_type() const
{
	return _eigen_problem_type ;
}
	
inline
bool NeutronicsSystem::is_eigen_problem_matrix_free() const
{
	return _is_eigen_problem_mfree ;
}


inline
bool& NeutronicsSystem::recompute_group_blocks()
{
	return _recompute_group_blocks ;
}

inline
bool NeutronicsSystem::recompute_group_blocks() const
{
	return _recompute_group_blocks ;
}

inline
const Problem& NeutronicsSystem::get_problem() const
{
	return *_problem ;
}

inline
Real& NeutronicsSystem::eigen_shift()
{
	return _eigen_shift ;
}

inline
Real NeutronicsSystem::eigen_shift() const
{
	return _eigen_shift ;
}


inline
bool& NeutronicsSystem::only_diagonal_loss_matrix()
{
	return _only_diagonal_loss_matrix ;
}

inline
bool NeutronicsSystem::only_diagonal_loss_matrix() const
{
	return _only_diagonal_loss_matrix ;
}

inline
bool& NeutronicsSystem::only_block_diagonal_loss_matrix()
{
	return _only_block_diagonal_loss_matrix ;
}

inline
bool NeutronicsSystem::only_block_diagonal_loss_matrix() const
{
	return _only_block_diagonal_loss_matrix ;
}

inline
PetscVector<Real>& NeutronicsSystem::group_vector(unsigned int ig) const
{
	libmesh_assert(ig < _egroups) ;
	
	return *_group_block_vectors[ig] ;
}

inline
PetscMatrix<Real>& NeutronicsSystem::group_matrix(unsigned int ig, unsigned int jg) const
{
	libmesh_assert(ig < _egroups) ;
	
	// if jg > ig, this could lead to seg-faults since we allocated only lower
	// triangular block matrices
	libmesh_assert(jg <= ig) ;
	
	return *_group_block_matrices[ig*_egroups+jg] ;
}

inline
KSP& NeutronicsSystem::loss_matrix_solver()
{
	return _loss_matrix_solver ;
}

inline
KSP& NeutronicsSystem::group_solver(unsigned int ig) const
{
	libmesh_assert(ig < _egroups) ;
	
	return _group_block_solvers[ig] ;
}

inline
const PetscVector<Real>& NeutronicsSystem::get_solution() const
{
	return dynamic_cast<PetscVector<Real>&>(*_system.solution) ;
}

inline
PetscVector<Real>& NeutronicsSystem::get_solution()
{
	return dynamic_cast<PetscVector<Real>&>(*_system.solution) ;
}

inline
const PetscVector<Real>& NeutronicsSystem::get_residual() const
{
	return dynamic_cast<PetscVector<Real>&>(*_system.rhs) ;
}

inline
PetscVector<Real>& NeutronicsSystem::get_residual()
{
	return dynamic_cast<PetscVector<Real>&>(*_system.rhs) ;
}


#endif	// define  __neutronics_system_h__
